# -*- coding: utf-8 -*-
import numpy as np
import random
import pickle
import pathlib
import sys

sys.path.append("..")  # 添加上级目录到系统路径，方便导入上级目录的模块
import config  # 导入配置文件

# np.random.seed(10)
# random.seed(10)
IMAGE_NUM = 155  # 设定图像（镜像）数量


class Trace:
    """
    该类用于解析 image_stats.csv 和 layer_stats.csv 文件，
    以构建镜像和层之间的关系。
    """

    def __init__(self):
        """
        初始化 Trace 类，读取镜像和层的信息
        """
        self.image_stats = dict()  # 存储镜像信息，key 为镜像名，value 为包含的相关信息
        self.layer_stats = dict()  # 存储层信息，key 为层的哈希值，value 为相关统计信息

        # 读取 image_stats.csv 文件
        with open("./data_task/image_stats.csv", "r") as f:
            cnt = 0
            for lines in f.readlines():
                image_item = lines[:-1].split(",")  # 解析 CSV 行数据
                self.image_stats[image_item[0]] = image_item[1:]
                self.image_stats[image_item[0]].append(cnt)  # 添加索引编号
                cnt += 1
                if cnt == IMAGE_NUM:
                    break  # 读取指定数量的镜像数据

        # 读取 layer_stats.csv 文件
        with open("./data_task/layer_stats.csv", "r") as f:
            cnt = 0
            for lines in f.readlines():
                layer_item = lines[:-1].split(",")
                if int(layer_item[1:-1][3]) > 0:  # 判断该层是否被某个镜像包含
                    is_layer_in_any_image = False
                    _image_list = []
                    for i in layer_item[-1].split("|"):
                        if i[8:] in self.image_stats:
                            _image_list.append(i[8:])
                            is_layer_in_any_image = True
                            if len(self.image_stats[i[8:]]) == 5:
                                self.image_stats[i[8:]].append(
                                    []
                                )  # 添加一个空列表存储层
                            self.image_stats[i[8:]][5].append(layer_item[0])
                    if is_layer_in_any_image:
                        self.layer_stats[layer_item[0]] = layer_item[1:-1]
                        self.layer_stats[layer_item[0]].append(cnt)
                        self.layer_stats[layer_item[0]].append(_image_list)
                        cnt += 1

        # 存储所有镜像的名称列表
        self.image_name_list = list(self.image_stats.keys())

    def generate_trace(
        self, task_number=20, task_rate=3, max_num_image=1, task_cpu_size_max=100
    ):
        """
        生成任务轨迹（泊松分布任务到达模式）
        :param task_number: 总任务数
        :param task_rate: 泊松分布的参数，控制任务到达速率
        :param max_num_image: 每个任务最多关联的镜像数
        :param task_cpu_size_max: 任务最大 CPU 需求
        """
        self.task_list = []
        self.task_sequence = []
        while sum(self.task_sequence) < task_number:
            self.task_sequence.append(
                np.random.poisson(task_rate)
            )  # 泊松分布生成任务数
            self.task_list.append([])

        task_number = sum(self.task_sequence)
        for i in range(len(self.task_list)):
            for _ in range(self.task_sequence[i]):
                image_reg_num = random.randint(1, max_num_image)
                self.task_list[i].append(
                    [
                        random.sample(
                            self.image_name_list, image_reg_num
                        ),  # 任务需要的镜像
                        random.random() * task_cpu_size_max,  # 随机 CPU 需求
                    ]
                )
        return self.task_list

    def generate_trace2(
        self,
        task_number=config.USER_NUM,
        arrival_time=0,
        max_time=config.TIME_ZONE,
        max_num_image=1,
        task_cpu_size_max=32,
        task_mem_size_max=20,
        task_average_interval=10,
        distribution="uniform",
        zipf_param=1.2,
        curseed=10,
    ):
        """
        生成具有随机到达时间的任务轨迹
        :param task_number: 任务总数
        :param arrival_time: 初始到达时间
        :param max_time: 任务最大时间范围
        :param max_num_image: 每个任务的最大镜像数
        :param task_cpu_size_max: 任务最大 CPU 需求
        :param task_mem_size_max: 任务最大内存需求
        :param task_average_interval: 任务到来的平均时间间隔
        :param distribution: 任务分布方式（uniform 或 zipf）
        :param zipf_param: zipf 分布的参数
        :param curseed: 随机种子
        """
        np.random.seed(curseed)
        random.seed(curseed)

        self.task_list = []
        if distribution == "uniform":
            for tid in range(task_number):
                image_reg_num = random.randint(1, max_num_image)
                image_name_list = self.image_name_list[:IMAGE_NUM]
                self.task_list.append(
                    [
                        tid,
                        round(random.uniform(arrival_time, max_time)),
                        random.sample(image_name_list, image_reg_num),
                        random.random() * task_cpu_size_max,
                        random.random() * task_mem_size_max,
                        list(
                            random.sample(
                                range(config.USER_NUM), config.task_belong_userNUM
                            )
                        ),
                    ]
                )
        return self.task_list

    def save_trace(self, file_name=None):
        """
        保存任务轨迹到文件
        """
        if file_name is not None:
            self.trace_file = file_name
        if self.task_list:
            with open(self.trace_file, "wb") as f:
                pickle.dump(self.task_list, f)
        else:
            print("No trace generated yet")

    def load_trace(self, file_name=None):
        """
        从文件加载任务轨迹
        """
        if file_name is not None:
            self.trace_file = file_name
        path = pathlib.Path(self.trace_file)
        if path.exists():
            with open(self.trace_file, "rb") as f:
                self.task_list = pickle.load(f)
        else:
            print("No such file")


if __name__ == "__main__":
    pass  # 主函数入口
